import glob
import logging
import os
import traceback

import hjson
import hvac

from absl import app, flags

FLAGS = flags.FLAGS

def main(_):
  src_path = FLAGS.src_path
  dst_path = FLAGS.dst_path
  vault_addr = FLAGS.vault_addr
  mount_point = FLAGS.mount_point
  assert os.path.exists(src_path), src_path
  assert dst_path is not None
  failed_keys = []
  client = hvac.Client(url=vault_addr)
  assert client.is_authenticated()
  for each_path in glob.iglob(f'{src_path}/**', recursive=True):
    try:
      if os.path.isfile(each_path) and each_path.endswith('.json'):
        with open(each_path) as fd:
          key_dict = hjson.load(fd)
        rel_path = os.path.relpath(each_path, start=src_path)
        vault_path = os.path.join(dst_path, rel_path)
        create_response = client.secrets.kv.v2.create_or_update_secret(
            path=vault_path, secret=key_dict, mount_point=mount_point)
    except Exception as e:
      logging.info('Fail to import: %s' % each_path)
      logging.info(traceback.format_exc())
      failed_keys.append(each_path)
  if len(failed_keys) > 0:
    logging.info('Fail to import keys: \n%s' % failed_keys)
  return len(failed_keys)


if __name__ == '__main__':
  logging.basicConfig(
      level='DEBUG', format='%(levelname)8s %(asctime)s %(name)s] %(message)s')
  flags.DEFINE_string('vault_addr', None, '')
  flags.DEFINE_string('src_path', None, '')
  flags.DEFINE_string('mount_point', None, '')
  flags.DEFINE_string('dst_path', None, '')

  app.run(main)
